import os
import requests
from datetime import datetime
from PIL import Image
from io import BytesIO
import time
import threading  # 添加线程模块

# 获取环境变量中的路径和 API Key
bozo = os.environ.get("BOZO") + "/Downloads/"
os.chdir(bozo)

# 定义尺寸映射字典，键为数字选项，值为对应的尺寸字符串
size_mapping = {
    1: "1536x2730",  # 9:16 比例（竖屏）- 1536*2731=4,194,816
    2: "2730x1536",  # 16:9 比例（横屏）- 2731*1536=4,194,816
    3: "2048x2048",  # 1:1 比例（正方形）- 2048*2048=4,194,304
    4: "2300x1800",  # 4:3 比例（传统显示器）- 2365*1774=4,194,310
    5: "1800x2300",  # 3:4 比例（竖向4:3）- 1774*2365=4,194,310
    6: "1600x2400",  # 2:3 比例（超宽屏）- 1672*2508=4,193,376
    7: "2400x1600",  # 3:2 比例（35mm胶片）- 2508*1672=4,193,376
    8: "2520x1080",  # 21:9 比例（超宽屏）- 3129*1341=4,196,989
}

# 1024尺寸映射字典
# size_mapping = {
# 1: "768x1365", # 9:16 比例（竖屏）
# 2: "1365x768", # 16:9 比例（横屏）
# 3: "1024x1024", # 1:1 比例（正方形）
# 4: "1152x910", # 4:3 比例（传统显示器）
# 5: "910x1152", # 3:4 比例（竖向4:3）
# 6: "836x1254", # 2:3 比例（超宽屏）
# 7: "1254x836", # 3:2 比例（35mm胶片）
# 8: "1566x670", # 21:9 比例（超宽屏）
# }

# 1280尺寸映射字典
# size_mapping = {
#     1: "960x1700",   # 9:16 比例（竖屏）- 960*1707=1,638,720
#     2: "1700x960",   # 16:9 比例（横屏）- 1707*960=1,638,720
#     3: "1280x1280",  # 1:1 比例（正方形）- 1280*1280=1,638,400
#     4: "1480x1100",  # 4:3 比例（传统显示器）- 1478*1109=1,639,102
#     5: "1100x1480",  # 3:4 比例（竖向4:3）- 1109*1478=1,639,102
#     6: "1040x1560",  # 2:3 比例（超宽屏）- 1045*1568=1,638,560
#     7: "1560x1040",  # 3:2 比例（35mm胶片）- 1568*1045=1,638,560
#     8: "1950x840",   # 21:9 比例（超宽屏）- 1955*838=1,638,290
# }

# 1536 尺寸
# size_mapping = {
#     1: "1152x2048",  # 9:16 比例（竖屏）
#     2: "2048x1152",  # 16:9 比例（横屏）
#     3: "1536x1536",  # 1:1 比例（正方形）
#     4: "1772x1329",  # 4:3 比例（传统显示器）
#     5: "1329x1772",  # 3:4 比例（竖向4:3）
#     6: "2352x1008",  # 2:3 比例（超宽屏）
#     7: "1881x1254",  # 3:2 比例（35mm胶片）
#     8: "2352x1008",  # 21:9 比例（超宽屏）
# }

# 获取用户输入的图片尺寸选项
size_option = int(input("请输入图片尺寸数字： (1➡️ 9:16、2➡️ 16:9、3➡️ 1:1、4➡️ 4:3、5➡️ 3:4、6➡️ 2:3、7➡️ 3:2、8➡️ 21:9)\r\n "))
# 根据用户选择获取对应的尺寸值
selected_size = size_mapping[size_option]

# 获取用户输入的中文提示词
tsc = input("请输入中文提示词：\r\n")

# 定义 API 请求的 URL
url = "https://ark.cn-beijing.volces.com/api/v3/images/generations"

api_key = os.environ.get("ARK_API_KEY")

# 定义请求头
headers = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {api_key}"
}

# 定义请求体
data = {
    "model": "doubao-seedream-3-0-t2i-250415",
    "prompt": tsc,
    "response_format": "url",
    "size": selected_size,
    "seed": -1,
    "guidance_scale": 2.5,
    "watermark": False
}

# 进度条控制变量
progress_running = True
progress_count = 0

# 进度条显示函数
def show_progress():
    global progress_count
    while progress_running:
        print("=", end="", flush=True)
        progress_count += 1
        time.sleep(1)

# 记录开始时间
start_time = time.time()

# 启动进度条线程
print("正在生成图片，请稍候: ", end="")
progress_thread = threading.Thread(target=show_progress)
progress_thread.start()

try:
    # 发送 POST 请求
    response = requests.post(url, headers=headers, json=data)
    
    # 检查请求是否成功
    if response.status_code == 200:
        # 解析响应 JSON 数据
        response_data = response.json()
        doubaoimg = response_data["data"][0]["url"]
        
        # 生成基于时间戳的文件名
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"豆包_{timestamp}.jpeg"
        
        # 下载图片并保存
        image = Image.open(BytesIO(requests.get(doubaoimg).content))
        image.save(filename)
        
        # 打开图片
        # os.system('open '+ bozo+filename)
        
        print(f"\n图片URL地址: {doubaoimg}")
    else:
        print(f"\n请求失败，状态码：{response.status_code}，错误信息：{response.text}")

finally:
    # 停止进度条线程
    progress_running = False
    progress_thread.join()
    
    # 计算并显示耗时
    end_time = time.time()
    elapsed_time = end_time - start_time
    print(f"\n图片生成完成! 总耗时: {elapsed_time:.2f} 秒")
    print(f"进度条显示时间: {progress_count} 秒")
